{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# TPU の使用\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/guide/tpu\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\"> TensorFlow.orgで表示</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Run in Google Colab</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">View source on GitHub</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">ノートブックをダウンロード/a0}</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ys81cOhXOWUP"
      },
      "source": [
        "Keras と Google Colab では現在、Cloud TPU を実験的にサポートしています。この Colab ノートブックをダウンロードする前に、Runtime > Change runtime type >  Hardware accelerator > TPU でノートブックの設定を確認し、ハードウェアアクセラレータが TPU であることを確認してください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ek5Hop74NVKm"
      },
      "source": [
        "## セットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cw0WRaChRxTL"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import os\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDWaRxSpwBN1"
      },
      "source": [
        "## TPU の初期化\n",
        "\n",
        "TPU は通常、ユーザーの Python プログラムを実行するローカルプロセスとは異なる Cloud TPU ワーカー上にあります。そのため、リモートクラスタに接続して TPU を初期化するには、ある程度の初期化作業が必要となります。`TPUClusterResolver` の `tpu` 引数は、Colab だけの特別なアドレスであることに注意してください。Google Compute Engine（GCE）で実行している場合は、ご利用の CloudTPU の名前を渡す必要があります。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dCqWMqvtwOLs"
      },
      "source": [
        "注意: TPU の初期化コードはプログラムのはじめにある必要があります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dKPqF8d1wJCV"
      },
      "outputs": [],
      "source": [
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])\n",
        "tf.config.experimental_connect_to_cluster(resolver)\n",
        "# This is the TPU initialization code that has to be at the beginning.\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "print(\"All devices: \", tf.config.list_logical_devices('TPU'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mv7kehTZ1Lq_"
      },
      "source": [
        "## 手動でデバイスを配置する\n",
        "\n",
        "TPU が初期されたら、計算を単一の TPU デバイスに配置するために、手動によるデバイスの配置を使用できます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRZ4kMoxBNND"
      },
      "outputs": [],
      "source": [
        "a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n",
        "b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
        "with tf.device('/TPU:0'):\n",
        "  c = tf.matmul(a, b)\n",
        "print(\"c device: \", c.device)\n",
        "print(c)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_NJm-kgFO0cC"
      },
      "source": [
        "## 分散ストラテジー\n",
        "\n",
        "ほとんどの場合、ユーザーはデータの並列化によって複数の TPU でモデルを実行することを望んでいます。分散ストラテジーは、モデルを CPU、GPU、または TPU で駆動するために使用できる抽象化です。分散ストラテジーを切り替えるだけで、モデルは特定のデバイスで実行するようになります。詳細は、「[分散ストラテジーガイド](./distributed_training.ipynb)」を参照してください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DcDPMZs-9uLJ"
      },
      "source": [
        "最初に、`TPUStrategy` オブジェクトを作成します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7SO23K8oRpjI"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.experimental.TPUStrategy(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JlaAmswWPsU6"
      },
      "source": [
        "計算を `strategy.run` API に渡すだけで、計算を複製してすべての TPU コアで実行できるようにすることができます。次のコード例では、すべてのコアが同じ入力 `(a, b)` を取得し、各コアで独立して matmul を実行しています。出力は、すべてのレプリカからの値となります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-90CL5uFPTOa"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def matmul_fn(x, y):\n",
        "  z = tf.matmul(x, y)\n",
        "  return z\n",
        "\n",
        "z = strategy.run(matmul_fn, args=(a, b))\n",
        "print(z)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uxgYl6kGHJLc"
      },
      "source": [
        "## TPU での分類\n",
        "\n",
        "基本的な概念を学んだので、より具体的な例を見てみることにしましょう。このガイドでは、分散ストラテジー `tf.distribute.experimental.TPUStrategy` を使用して Cloud TPU を駆動し、Keras モデルをトレーニングする方法を説明します。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKRALGgt_kCo"
      },
      "source": [
        "### Keras モデルを定義する\n",
        "\n",
        "次のコードは、Keras を使用した MNIST モデルの定義です。CPU または GPU で使用されるものと変わりません。変数が各 TPU デバイスに作成されるように、Keras モデルの作成を `strategy.scope` 内で行う必要があることに注意してください。ほかの部分のコードは、ストラテジーのスコープ内でなくても構いません。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DiBiN-Z_R7P7"
      },
      "outputs": [],
      "source": [
        "def create_model():\n",
        "  return tf.keras.Sequential(\n",
        "      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),\n",
        "       tf.keras.layers.Conv2D(256, 3, activation='relu'),\n",
        "       tf.keras.layers.Flatten(),\n",
        "       tf.keras.layers.Dense(256, activation='relu'),\n",
        "       tf.keras.layers.Dense(128, activation='relu'),\n",
        "       tf.keras.layers.Dense(10)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qYOYjYTg_31l"
      },
      "source": [
        "### 入力データセット\n",
        "\n",
        "Cloud TPU にデータを迅速にフィードできなければ Cloud TPU を使用することは不可能であるため、Cloud TPU を使用する際は、`tf.data.Dataset` API を効率的に使用できることが非常に重要となります。データセットのパフォーマンスに関する詳細は、ｐ[入力パイプラインのパフォーマンスガイド](./data_performance.ipynb)」を参照してください。\n",
        "\n",
        "最も単純な実験（`tf.data.Dataset.from_tensor_slices` またはほかのグラフ内データの使用）以外のすべての実験では、データセットが読み取るすべてのデータファイルを Google Cloud Storage（GCS）バケットに格納する必要があります。\n",
        "\n",
        "ほとんどの使用事例では、データを `TFRecord` 形式に変換し、`tf.data.TFRecordDataset` を使って読み取ることをお勧めします。このやり方については、「[TFRecord および tf.Example のチュートリアル](../tutorials/load_data/tfrecord.ipynb)」を参照してください。ただし、これは絶対要件ではないため、ほかのデータセットリーダー（`FixedLengthRecordDataset` または `TextLineDataset`）を使用することもできます。\n",
        "\n",
        "小さなデータセットは、`tf.data.Dataset.cache` を使ってすべてをメモリに読み込むことができます。\n",
        "\n",
        "データ形式にかかわらず、100 MB 程度の大きなファイルを使用することをお勧めします。このネットワーク化された設定においては、ファイルを開くタスクのオーバーヘッドが著しく高くなるため、特に重要なことです。\n",
        "\n",
        "ここでは、`tensorflow_datasets` モジュールをし米須ヒデ、MNIST トレーニングデータのコピーを取得する必要があります。`try_gcs` は、パブリック GCS バケットで提供されているコピーを使用するように指定されています。 これを指定しない場合、TPUはダウンロードされたデータにアクセスできません。 "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "noAd416KSCo7"
      },
      "outputs": [],
      "source": [
        "def get_dataset(batch_size, is_training=True):\n",
        "  split = 'train' if is_training else 'test'\n",
        "  dataset, info = tfds.load(name='mnist', split=split, with_info=True,\n",
        "                            as_supervised=True, try_gcs=True)\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255.0\n",
        "\n",
        "    return image, label\n",
        "\n",
        "  dataset = dataset.map(scale)\n",
        "\n",
        "  # Only shuffle and repeat the dataset in training. The advantage to have a\n",
        "  # infinite dataset for training is to avoid the potential last partial batch\n",
        "  # in each epoch, so users don't need to think about scaling the gradients\n",
        "  # based on the actual batch size.\n",
        "  if is_training:\n",
        "    dataset = dataset.shuffle(10000)\n",
        "    dataset = dataset.repeat()\n",
        "\n",
        "  dataset = dataset.batch(batch_size)\n",
        "\n",
        "  return dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mgUC6A-zCMEr"
      },
      "source": [
        "### Keras の高位 API を使用してモデルをトレーニングする\n",
        "\n",
        "Keras の fit/compile API を使用するだけでモデルをトレーニングできます。ここでは、TPU 固有のものはないため、複数の GPU を使用している場合や、`TPUStrategy` ではなく`MirroredStrategy` を使用している場合でも、以下のコードを記述できます。詳細については、「[Kera wo使用した分散トレーニング](https://www.tensorflow.org/tutorials/distribute/keras)」チュートリアルを参照してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubmDchPqSIx0"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  model.compile(optimizer='adam',\n",
        "                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                metrics=['sparse_categorical_accuracy'])\n",
        "\n",
        "batch_size = 200\n",
        "steps_per_epoch = 60000 // batch_size\n",
        "validation_steps = 10000 // batch_size\n",
        "\n",
        "train_dataset = get_dataset(batch_size, is_training=True)\n",
        "test_dataset = get_dataset(batch_size, is_training=False)\n",
        "\n",
        "model.fit(train_dataset,\n",
        "          epochs=5,\n",
        "          steps_per_epoch=steps_per_epoch,\n",
        "          validation_data=test_dataset, \n",
        "          validation_steps=validation_steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8hSGBIYtUugJ"
      },
      "source": [
        "Python のオーバーヘッドを緩和し、TPU のパフォーマンスを最大化するには、`Model.compile` の **実験的な** `experimental_steps_per_execution` 引数を使用してみてください。スループットが約 50% 改善されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M6e3aVVLUorL"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  model.compile(optimizer='adam',\n",
        "                # Anything between 2 and `steps_per_epoch` could help here.\n",
        "                experimental_steps_per_execution = 50,\n",
        "                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                metrics=['sparse_categorical_accuracy'])\n",
        "\n",
        "model.fit(train_dataset,\n",
        "          epochs=5,\n",
        "          steps_per_epoch=steps_per_epoch,\n",
        "          validation_data=test_dataset,\n",
        "          validation_steps=validation_steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0rRALBZNCO4A"
      },
      "source": [
        "### カスタムトレーニングループを使用してモデルをトレーニングする\n",
        "\n",
        "`tf.function` と `tf.distribute` API を直接使用しても、モデルを作成し、トレーニングすることができます。`strategy.experimental_distribute_datasets_from_function` API は、データセット関数を指定してデータセットを配布するために使用されます。この場合、データセットに渡されるバッチサイズは、グローバルバッチサイズではなく、レプリカごとのバッチサイズであることに注意してください。詳細については、「[tf.distribute.Strategy によるカスタムトレーニング](https://www.tensorflow.org/tutorials/distribute/custom_training)」チュートリアルを参照してください。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DxdgXPAL6iFE"
      },
      "source": [
        "最初に、モデル、データセット、および tf.functions を作成します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9aHhqwao2Fxi"
      },
      "outputs": [],
      "source": [
        "# Create the model, optimizer and metrics inside strategy scope, so that the\n",
        "# variables can be mirrored on each device.\n",
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  optimizer = tf.keras.optimizers.Adam()\n",
        "  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)\n",
        "  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "      'training_accuracy', dtype=tf.float32)\n",
        "\n",
        "# Calculate per replica batch size, and distribute the datasets on each TPU\n",
        "# worker.\n",
        "per_replica_batch_size = batch_size // strategy.num_replicas_in_sync\n",
        "\n",
        "train_dataset = strategy.experimental_distribute_datasets_from_function(\n",
        "    lambda _: get_dataset(per_replica_batch_size, is_training=True))\n",
        "\n",
        "@tf.function\n",
        "def train_step(iterator):\n",
        "  \"\"\"The step function for one training step\"\"\"\n",
        "\n",
        "  def step_fn(inputs):\n",
        "    \"\"\"The computation to run on each TPU device.\"\"\"\n",
        "    images, labels = inputs\n",
        "    with tf.GradientTape() as tape:\n",
        "      logits = model(images, training=True)\n",
        "      loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
        "          labels, logits, from_logits=True)\n",
        "      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)\n",
        "    grads = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
        "    training_loss.update_state(loss * strategy.num_replicas_in_sync)\n",
        "    training_accuracy.update_state(labels, logits)\n",
        "\n",
        "  strategy.run(step_fn, args=(next(iterator),))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ibi7Z97V6xsQ"
      },
      "source": [
        "次に、トレーニングループを実行します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1du5cXWt6Vtw"
      },
      "outputs": [],
      "source": [
        "steps_per_eval = 10000 // batch_size\n",
        "\n",
        "train_iterator = iter(train_dataset)\n",
        "for epoch in range(5):\n",
        "  print('Epoch: {}/5'.format(epoch))\n",
        "\n",
        "  for step in range(steps_per_epoch):\n",
        "    train_step(train_iterator)\n",
        "  print('Current step: {}, training loss: {}, accuracy: {}%'.format(\n",
        "      optimizer.iterations.numpy(),\n",
        "      round(float(training_loss.result()), 4),\n",
        "      round(float(training_accuracy.result()) * 100, 2)))\n",
        "  training_loss.reset_states()\n",
        "  training_accuracy.reset_states()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TnZJUM3qIjKu"
      },
      "source": [
        "### `tf.function` 内の複数のステップでパフォーマンスを改善する\n",
        "\n",
        "`tf.function` 内で複数のステップを実行することで、パフォーマンスを改善できます。これは、`tf.function` 内の `tf.range` で `strategy.run` 呼び出しをラッピングすることで実現されます。AutoGraph は、TPU ワーカーの `tf.while_loop` に変換します。\n",
        "\n",
        "パフォーマンスは改善されますが、`tf.function` 内の単一のステップに比べれば、トレードオフがあります。`tf.function` で複数のステップを実行すると柔軟性に劣り、ステップ内での Eager execution や任意の Python コードを実行できません。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2grYvXLzJYkP"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_multiple_steps(iterator, steps):\n",
        "  \"\"\"The step function for one training step\"\"\"\n",
        "\n",
        "  def step_fn(inputs):\n",
        "    \"\"\"The computation to run on each TPU device.\"\"\"\n",
        "    images, labels = inputs\n",
        "    with tf.GradientTape() as tape:\n",
        "      logits = model(images, training=True)\n",
        "      loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
        "          labels, logits, from_logits=True)\n",
        "      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)\n",
        "    grads = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
        "    training_loss.update_state(loss * strategy.num_replicas_in_sync)\n",
        "    training_accuracy.update_state(labels, logits)\n",
        "\n",
        "  for _ in tf.range(steps):\n",
        "    strategy.run(step_fn, args=(next(iterator),))\n",
        "\n",
        "# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get \n",
        "# retraced if the value changes.\n",
        "train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))\n",
        "\n",
        "print('Current step: {}, training loss: {}, accuracy: {}%'.format(\n",
        "      optimizer.iterations.numpy(),\n",
        "      round(float(training_loss.result()), 4),\n",
        "      round(float(training_accuracy.result()) * 100, 2)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WBKVhMvWjibf"
      },
      "source": [
        "## 次のステップ\n",
        "\n",
        "- [Google Cloud TPU ドキュメント](https://cloud.google.com/tpu/docs/) - Google Cloud TPU をセットアップして実行します。\n",
        "- [TensorFlow での分散型トレーニング](./distributed_training.ipynb) - 分散ストラテジーの使用方法とベストプラクティスを示す多数の例へのリンクが含まれます。\n",
        "- [TensorFlow でモデルを保存/読み込む](https://github.com/tensorflow/models/tree/master/official) - 分散ストラテジーでモデルを保存および読み込む方法を説明します。\n",
        "- [TensorFlow 公式モデル](https://cloud.google.com/tpu/docs/performance-guide) - Cloud TPU 対応の最新の TensorFlow 2.x モデルの例を紹介します。\n",
        "- [Google Cloud TPU パフォーマンスガイド](https://cloud.google.com/tpu/docs/performance-guide) - アプリケーションに合った Cloud TPU 構成パラメータの調整し、Cloud TPU パフォーマンスをさらに改善します。"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "tpu.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
